import os
import csv
import pymysql
from config import MYSQL_CONFIG


# 判断文件是否被其他程序占用（用于跳过）
def is_file_in_use(file_path):
    try:
        with open(file_path, 'r', encoding='utf-8'):
            return False
    except IOError:
        return True


def get_db_connection():
    """获取数据库连接"""
    return pymysql.connect(**MYSQL_CONFIG)


def add_column_if_not_exists(cursor, table_name, column_name):
    """如果列不存在，则添加新列"""
    try:
        cursor.execute(f"SHOW COLUMNS FROM `{table_name}` LIKE '{column_name}'")
        result = cursor.fetchone()
        if not result:
            alter_sql = f"ALTER TABLE `{table_name}` ADD COLUMN `{column_name}` TEXT;"
            cursor.execute(alter_sql)
            print(f"已新增列: `{column_name}`")
        else:
            print(f"列 `{column_name}` 已存在，跳过添加")
    except Exception as e:
        print(f"添加列 `{column_name}` 失败: {e}")


def create_table_if_not_exists(cursor, table_name, columns):
    """如果表不存在，则创建表，并包含所有列"""
    try:
        cursor.execute(f"SHOW TABLES LIKE '{table_name}'")
        result = cursor.fetchone()
        if not result:
            columns_sql = ", ".join([f"`{col}` TEXT" for col in columns])
            create_table_sql = f"CREATE TABLE `{table_name}` (id INT AUTO_INCREMENT PRIMARY KEY, {columns_sql});"
            cursor.execute(create_table_sql)
            print(f"已创建表: `{table_name}`")
        else:
            print(f"表 `{table_name}` 已存在，跳过创建")
    except Exception as e:
        print(f"创建表失败: {e}")


def insert_csv_to_mysql(csv_path, table_name):
    """将单个 CSV 文件插入数据库"""
    if is_file_in_use(csv_path):
        print(f"文件被占用，跳过: {csv_path}")
        return

    try:
        with open(csv_path, 'r', encoding='utf-8') as file:
            reader = csv.reader(file)
            headers = next(reader)  # 获取 CSV 表头

            conn = get_db_connection()
            cursor = conn.cursor()

            # 检查表是否存在并创建
            create_table_if_not_exists(cursor, table_name, headers)

            # 确保每个列都存在（用于后续增量更新）
            for header in headers:
                add_column_if_not_exists(cursor, table_name, header)

            # 插入数据
            for row in reader:
                placeholders = ', '.join(['%s'] * len(row))
                insert_sql = f"INSERT INTO `{table_name}` ({', '.join([f'`{h}`' for h in headers])}) VALUES ({placeholders})"
                cursor.execute(insert_sql, row)

            conn.commit()  # 提交事务
            print(f"已成功导入文件: {csv_path}")

    except Exception as e:
        print(f" 处理文件时出错: {e}, 文件: {csv_path}")
        if 'conn' in locals():
            conn.rollback()
    finally:
        if 'cursor' in locals():
            cursor.close()
        if 'conn' in locals():
            conn.close()


def process_directory(directory, table_name):
    """遍历目录，处理所有 part-00000 开头的 CSV 文件"""
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.startswith('part-00000'):
                csv_path = os.path.join(root, file)
                insert_csv_to_mysql(csv_path, table_name)


if __name__ == "__main__":
    from config import LOCAL_DIR  # 从 config 导入路径

    directory = LOCAL_DIR  # 使用 config 中定义的本地保存路径
    table_name = "transport_task_final"

    process_directory(directory, table_name)